## ! /usr/bin/python
# -*- coding: utf-8 -*-
# @author HanYZ
# @createAt 2022.07.17

from os.path import expanduser, dirname, exists, getmtime, join
from dateutil.relativedelta import relativedelta
from datetime import datetime as dt
from random import randint
from os import makedirs
import argparse
import baostock as bs
import pandas as pd
import sys

_datadir = join(expanduser('~'), '.local', 'datas')
_sharedatadir = '/usr/local/share/datas' if sys.platform == 'linux' else _datadir

_config = dict(
    day=dict(
        fields="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
        filepath=lambda code: join(_datadir, 'day', code[:6], code+'.csv'),
        sharefilepath=lambda code: join(
            _sharedatadir, 'day', code[:6], code+'.csv'),
        listpath=join(_datadir, 'day', 'all_stock.csv'),
        frequency='d',
    ),
    week=dict(
        fields="date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg",
        filepath=lambda code: join(_datadir, 'week', code[:6], code+'.csv'),
        sharefilepath=lambda code: join(
            _sharedatadir, 'week', code[:6], code+'.csv'),
        listpath=join(_datadir, 'week', 'all_stock.csv'),
        frequency='w',
    ),
    month=dict(
        fields="date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg",
        filepath=lambda code: join(_datadir, 'month', code[:6], code+'.csv'),
        sharefilepath=lambda code: join(
            _sharedatadir, 'month', code[:6], code+'.csv'),
        listpath=join(_datadir, 'month', 'all_stock.csv'),
        frequency='m',
    ),
)

_cache = dict(d=dict(), w=dict(), m=dict())


def _lastDay(conf):
    now = dt.now()
    ld = now + relativedelta(hour=17, minute=35)
    if conf['frequency'] == 'm':
        ld += relativedelta(months=1, day=1, days=-1)
        if ld.weekday() > 4:
            ld += relativedelta(days=4-ld.weekday())
    if ld > now:
        if conf['frequency'] == 'm':
            ld += relativedelta(day=1, days=-1)
        else:
            ld += relativedelta(days=-1)
    if ld.weekday() > 4:
        ld += relativedelta(days=4-ld.weekday())
    elif ld.weekday() < 4 and conf['frequency'] == 'w':
        ld += relativedelta(days=-3-ld.weekday())
    return ld


def _read_csv(filepath, klen, skippause=False):
    data = pd.read_csv(filepath, index_col='date', parse_dates=True)
    if skippause and ('tradestatus' in data):
        data = data[data.tradestatus > 0]
    if klen > 0:
        data = data.tail(klen)
    return data


def _to_csv(dataframe, filepath):
    d = dirname(filepath)
    if not exists(d):
        makedirs(d)
        print('创建', d)
    dataframe.to_csv(filepath, index=0)


def _downloadOne(code, conf, start_date='1991-01-01', write=True):
    end_date = dt.now().date().isoformat()
    rs = bs.query_history_k_data_plus(code, conf['fields'], start_date=start_date, end_date=end_date,
                                      frequency=conf['frequency'], adjustflag='2').get_data()

    # 写入文件
    if len(rs) > 0:
        write and _to_csv(rs, conf['filepath'](code))
    else:
        print('%s 的数据未读到' % code)
    return rs


def _refresh(code, conf, login=True):
    filepath = conf['filepath'](code)
    sharefilepath = conf['sharefilepath'](code)

    if exists(sharefilepath) and dt.fromtimestamp(getmtime(sharefilepath)) > _lastDay(conf):  # 共享数据是新的
        return False

    if exists(filepath):  # 数据文件存在
        if dt.fromtimestamp(getmtime(filepath)) > _lastDay(conf):  # 数据是新的
            return False
        if login:
            bs.login()
        old = _read_csv(filepath, 3, False)
        if len(old) == 3:
            last = old.index[0]
            rs = _downloadOne(
                code, conf, start_date=last.date().isoformat(), write=False)
            # 没有复权等情况影响历史数据
            if len(rs) > 2 and old.open[2] == rs['open'].astype(float)[2]:
                rs.loc[3:, :].to_csv(filepath, mode='a',
                                     index=False, header=None)
                if login:
                    bs.logout()
                return True
    elif login:
        bs.login()

    _downloadOne(code, conf)  # 重新下载
    bs.logout
    return True


def _kdata(code, klen, conf, download=False, skippause=True):
    cache = _cache[conf['frequency']]
    if code not in cache:
        filepath = conf['filepath'](code)
        sharefilepath = conf['sharefilepath'](code)
        if not exists(sharefilepath) and not exists(filepath):
            download = True
        download and _refresh(code, conf, True)
        if exists(sharefilepath) and (not exists(filepath) or getmtime(sharefilepath) > getmtime(filepath)):
            filepath = sharefilepath
        cache[code] = _read_csv(filepath, 0, False)
    data = cache[code]
    if skippause and ('tradestatus' in data):
        data = data[data.tradestatus > 0]
    if klen > 0:
        data = data.tail(klen)

    return data.copy()


def _download(conf):
    bs.login()
    if conf['frequency'] == 'w':
        goodgoods = dict(
            sz50=bs.query_sz50_stocks,
            hs300=bs.query_hs300_stocks,
            zz500=bs.query_zz500_stocks,
        )
        for name in goodgoods:
            data = goodgoods[name]().get_data()
            if len(data) == 0:
                continue
            data.to_csv(join(_datadir, f'{name}_stock.csv'))
    date = _lastDay(conf)
    while True:
        all_stock = bs.query_all_stock(date.date().isoformat()).get_data()
        if len(all_stock) > 0:
            break
        date += relativedelta(days=-1)

    cs = all_stock[~(all_stock.code.str.startswith('bj') | (all_stock.tradeStatus == '0') |
                     all_stock.code_name.str.contains('退') | all_stock.code_name.str.contains('ST'))]
    _to_csv(cs, conf['listpath'])

    for i in cs.index:
        c = cs.loc[i]
        _refresh(cs.loc[i].code, conf, login=False)

    bs.logout()
    pass


def _all_stock(conf):
    filepath = conf['listpath']
    if not exists(filepath):
        _download(conf)
    return pd.read_csv(filepath)


def daily(code, klen=0, download=False):
    return _kdata(code, klen, _config['day'], download)


def weekly(code, klen=0, download=False):
    return _kdata(code, klen, _config['week'], download)


def anyone(frequency='day', klen=10):
    '''任选一只证券，返回(Kdata, code, name)'''
    conf = _config[frequency]
    al = _all_stock(conf)
    while True:
        c = al.iloc[randint(0, len(al))]
        d = _kdata(c.code, klen, conf, False)
        if klen < 1 or len(d) >= klen:
            return d, c.code, c.code_name


def onebyone(frequency='day', klen=0, types=['sh.60', 'sz.00'], limitgt=0, start=0, end=0, withdata=True):
    '''遍历证券信息，每个包含(Kdata, code, name)'''
    conf = _config[frequency]
    al = _all_stock(conf)
    if limitgt > 0:
        klen = 0
    else:
        limitgt = klen
    count = 0
    for i in al.index:
        c = al.loc[i]
        if types and c.code[:5] not in types:
            continue
        count += 1
        if end > 0 and count >= end:
            break
        if count < start:
            continue
        if withdata:
            d = _kdata(c.code, klen, conf, False)
            if len(d) > 0 and (limitgt == 0 or len(d) >= limitgt):
                yield d, c.code, c.code_name
        else:
            yield None, c.code, c.code_name



def download(frequency='day'):
    '''下载证券数据'''
    print(dt.now(), '开始下载')
    _download(_config[frequency])
    print(dt.now(), '下载完成')


def goodcheckerbuilder():
    res = {}
    cname = dict(sz50='上证50', hs300='沪深300', zz500='中证500')
    for name in cname:
        fn = join(_datadir, f'{name}_stock.csv')
        if exists(fn):
            res[name] = pd.read_csv(fn)

    def foo(code):
        def check(c, d): return len(d[d.code == c]) > 0
        return [[check(code, res[name]), name, cname[name]] for name in res]
    return foo


def goodchecklist():
    '''获取好股票的(检查器列表,编码列表,名称列表)'''
    res = {}
    cname = dict(sz50='上证50', hs300='沪深300', zz500='中证500')

    for name in cname:
        fn = join(_datadir, f'{name}_stock.csv')
        if exists(fn):
            res[name] = pd.read_csv(fn)

    def check(c, d): return len(d[d.code == c]) > 0
    funcs = [(lambda d: check(d.code[0], res[k])) for k in cname]

    return funcs, list(cname.keys()), [cname[k] for k in cname]


def parse_args():
    parser = argparse.ArgumentParser(description='从baostock.com下载数')

    parser.add_argument('--frequency', default='day', required=False,
                        choices=['day', 'week', 'month'],
                        help='选择下载日线还是周线、月线，默认值：日线')

    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    download(args.frequency)
